# -*- coding: utf-8 -*-
# @Time    : 2019/7/23 17:32
# @Author  : zhouqiang
# @Site    : 
# @File    : portfolio_tool.py
# @Software: PyCharm

"""
脚本说明：组合相关的api函数，
之前是通过portfolio_interface中的函数调用易鑫那边的接口，然后将数据转换成DataFrame输出。
现在是从库里直接取数据
"""

import pandas as pd

from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.datasource_fetch.portfolio_api.portfolio_interface import get_portfolio_info_with_id
from quant_researcher.quant.project_tool.db_operator import db_conn


# def get_portfolio_basic_info(portfolio_id):
#     """
#     根据组合ID找到组合的基本信息
#
#     :param portfolio_id: 组合ID
#     :return: pd.DataFrame
#     """
#     res = get_portfolio_info_with_id(406009, {'combination_id': portfolio_id})
#     if res is None:
#         LOG.error(f"{portfolio_id} 未找到该组合的基本信息")
#         return
#     else:
#         return res


def get_portfolio_basic_info(portfolio_id):
    """
    根据组合ID找到组合的基本信息

    :param portfolio_id: 组合ID
    :return: pd.DataFrame
    """
    conn = db_conn.get_fof_conn()
    res = pd.read_sql(f"select tid as comb_id, user_id, combination_type, combination_name, combination_source, "
                      f"compref, is_valid, init_capital, inception_date, combination_biref, combination_describe, "
                      f"comb_risk, create_type "
                      f"from t_researchr_combinfo "
                      f"where tid = '{portfolio_id}' ", conn)
    conn.close()
    if res.empty:
        LOG.error(f"未找到{portfolio_id}该组合的基本信息")
        return
    else:
        return res


# def get_portfolio_nav_related_info(portfolio_id, start_date, end_date=None):
#     """
#     根据组合ID找到某一天或者某个时间段的组合净值相关数据，结果按照日期从小到大排序
#
#     :param portfolio_id: 组合ID
#     :param start_date: 开始时间，格式"20190701"
#     :param end_date: 结束时间，格式"20200701"，若取某一天的数据，end_date参数可不传，也可与start_date传同一天
#     :return: pd.DataFrame
#     """
#     if end_date is None:
#         end_date = start_date
#     nav_p = get_portfolio_info_with_id(403403,
#                                        {'combination_id': portfolio_id,
#                                         'start_date': start_date,
#                                         'end_date': end_date
#                                         },
#                                        paging=False)
#     if nav_p is None:
#         LOG.error(f"{portfolio_id} 未找到该组合的净值信息")
#         return
#     else:
#         nav_p = pd.DataFrame(data=nav_p)
#         nav_p = nav_p.where(nav_p != '')
#         str_list = ['account_id', 'trade_date']
#         nav_p[nav_p.columns.difference(str_list)] = nav_p[nav_p.columns.difference(str_list)].astype(float)
#         nav_p = nav_p.sort_values(by='trade_date')
#         return nav_p


def get_portfolio_nav_related_info(portfolio_id, start_date, end_date=None):
    """
    根据组合ID找到某一天或者某个时间段的组合净值相关数据，结果按照日期从小到大排序

    :param portfolio_id: 组合ID
    :param start_date: 开始时间，格式"20190701"
    :param end_date: 结束时间，格式"20200701"，若取某一天的数据，end_date参数可不传，也可与start_date传同一天
    :return: pd.DataFrame
    """
    if end_date is None:
        end_date = start_date

    conn = db_conn.get_fof_conn()
    nav_p = pd.read_sql(f"select nav_date as trade_date, nav as account_net_value "
                        f"from t_researchr_combnav "
                        f"where comb_id = '{portfolio_id}' "
                        f"and nav_date >= '{start_date}' "
                        f"and nav_date <= '{end_date}' "
                        f"order by nav_date", conn)
    conn.close()
    if nav_p.empty:
        LOG.error(f"未找到{portfolio_id}该组合的净值信息")
        return
    else:
        nav_p['trade_date'] = nav_p['trade_date'].astype(str)
        nav_p['trade_date'] = nav_p['trade_date'].apply(lambda x: x.replace('-', ''))
        return nav_p


# def get_portfolio_hold_related_info(portfolio_id, start_date, end_date=None):
#     """
#     根据组合ID找到某一天或者某个时间段的组合持仓相关数据
#
#     :param portfolio_id: 组合ID
#     :param start_date: 开始时间，格式"20190701"
#     :param end_date: 结束时间，格式"20200701"，若取某一天的数据，end_date参数可不传，也可与start_date传同一天
#     :return:
#     """
#     # 如果取一天的数据
#     if (end_date is None) or (start_date == end_date):
#         df = get_portfolio_info_with_id(403404,
#                                         {'combination_id': portfolio_id, 'end_date': start_date},
#                                         paging=False)
#     # 如果取一段时间的数据
#     else:
#         df = get_portfolio_info_with_id(403404,
#                                         {'combination_id': portfolio_id,
#                                          'start_date': start_date, 'end_date': end_date,
#                                          'is_total': '1'},
#                                         paging=False)
#     if df is None:
#         LOG.error(f"未找到{portfolio_id}组合的持仓信息")
#         return
#     else:
#         df = pd.DataFrame(data=df)
#         df = df.where(df != '')
#         str_list = ['hold_no', 'account_id', 'fund_code', 'fund_name', 'bonus_share',
#                     'create_date', 'modify_date', 'backup_date']
#         df[df.columns.difference(str_list)] = df[df.columns.difference(str_list)].astype(float)
#         return df


def get_portfolio_hold_related_info(portfolio_id, start_date, end_date=None):
    """
    根据组合ID找到某一天或者某个时间段的组合持仓相关数据

    :param portfolio_id: 组合ID
    :param start_date: 开始时间，格式"20190701"
    :param end_date: 结束时间，格式"20200701"，若取某一天的数据，end_date参数可不传，也可与start_date传同一天
    :return:
    """
    conn = db_conn.get_fof_conn()
    # 如果取一天的数据
    if (end_date is None) or (start_date == end_date):
        df = pd.read_sql(f"select backup_date, sec_code, sec_type, sec_name, "
                         f"total_qty, today_price, sec_balance, cost_price, cost_balance "
                         f"from t_researchr_combholding "
                         f"where comb_id = '{portfolio_id}' "
                         f"and backup_date = ("
                         f"select max(backup_date) "
                         f"from t_researchr_combholding "
                         f"where comb_id = '{portfolio_id}' "
                         f"and backup_date <= '{start_date}')", conn)
        sell_buy_df = pd.read_sql(f"select exec_date as backup_date, sec_code, sec_type, trade_type, bargain_balance "
                                  f"from t_researchr_combexecagge "
                                  f"where comb_id = '{portfolio_id}' "
                                  f"and exec_date = ("
                                  f"select max(exec_date) "
                                  f"from t_researchr_combexecagge "
                                  f"where comb_id = '{portfolio_id}' "
                                  f"and exec_date <= '{start_date}')", conn)
        bonus_df = pd.read_sql(f"select bonus_date as backup_date, sec_code, sec_type, bonus_money as bonus_share "
                               f"from t_researchr_combbonus "
                               f"where comb_id = '{portfolio_id}' "
                               f"and bonus_date = ("
                               f"select max(bonus_date) "
                               f"from t_researchr_combbonus "
                               f"where comb_id = '{portfolio_id}' "
                               f"and bonus_date <= '{start_date}')", conn)

    # 如果取一段时间的数据
    else:
        df = pd.read_sql(f"select backup_date, sec_code, sec_type, sec_name, "
                         f"total_qty, today_price, sec_balance, cost_price, cost_balance "
                         f"from t_researchr_combholding "
                         f"where comb_id = '{portfolio_id}' "
                         f"and backup_date >= '{start_date}' "
                         f"and backup_date <= '{end_date}' ", conn)
        sell_buy_df = pd.read_sql(f"select exec_date as backup_date, sec_code, sec_type, trade_type, bargain_balance "
                                  f"from t_researchr_combexecagge "
                                  f"where comb_id = '{portfolio_id}' "
                                  f"and exec_date >= '{start_date}' "
                                  f"and exec_date <= '{end_date}' ", conn)
        bonus_df = pd.read_sql(f"select bonus_date as backup_date, sec_code, sec_type, bonus_money as bonus_share "
                               f"from t_researchr_combbonus "
                               f"where comb_id = '{portfolio_id}' "
                               f"and bonus_date >= '{start_date}' "
                               f"and bonus_date <= '{end_date}' ", conn)
    conn.close()
    if df.empty:
        LOG.error(f"未找到{portfolio_id}组合的持仓信息")
        return
    else:
        sell_buy_df['sell_balance'] = sell_buy_df['bargain_balance'].where(sell_buy_df['trade_type'] == '31')
        sell_buy_df['buy_balance'] = sell_buy_df['bargain_balance'].where(sell_buy_df['trade_type'] == '30')
        sell_buy_df = sell_buy_df.drop(columns=['trade_type', 'bargain_balance'])

        df = df.merge(sell_buy_df, how='left', on=['backup_date', 'sec_code', 'sec_type'])
        df = df.merge(bonus_df, how='left', on=['backup_date', 'sec_code', 'sec_type'])
        df = df[df['sec_type'] == '1']
        df = df.rename(columns={'sec_code': 'fund_code', 'sec_name': 'fund_name',
                                'sec_balance': 'fund_balance', 'today_price': 'net_value'})
        df = df.drop(columns=['sec_type'])

        df['backup_date'] = df['backup_date'].astype(str)
        df['backup_date'] = df['backup_date'].apply(lambda x: x.replace('-', ''))

        return df


def get_portfolio_fund_weight(portfolio_id, start_date, end_date=None):
    """
    根据组合ID得到某一天或者某段时间的子基金权重

    :param portfolio_id:  str，组合ID
    :param start_date: 开始时间，格式"20190701"
    :param end_date: 结束时间，格式"20200701"，若取某一天的数据，end_date参数可不传，也可与start_date传同一天
    :return: pd.DataFrame
    """
    df = get_portfolio_hold_related_info(portfolio_id, start_date, end_date)
    if df is None:
        LOG.error(f"未找到{portfolio_id}组合的基金信息")
        return
    # 如果取一天的数据
    if (end_date is None) or (start_date == end_date):
        df = df[['fund_code', 'fund_balance']]
        df['fund_weight'] = df['fund_balance'] / df['fund_balance'].sum()
        df = df[['fund_code', 'fund_weight']]
    # 如果取一段时间的数据
    else:
        df = df.rename(columns={'backup_date': 'tj', 'fund_balance': 'value'})
        df = df.set_index(['tj', 'fund_code'])['value'].unstack()
        df = df.fillna(0)
        df = df.div(df.sum(axis=1), axis=0)
    return df


if __name__ == '__main__':
    # aaa = get_portfolio_fund_weight('1296', '20150130', '20210126')
    # aaa = get_portfolio_fund_weight('289', '20201201', '20201215')
    # bbb = get_portfolio_basic_info('1047')
    # ccc = get_portfolio_nav_related_info('1047', '20200301', '20210301')
    ddd = get_portfolio_fund_weight('1189', '20200701', '20200901')
    fff = get_portfolio_fund_weight('1189', '20200701')
